#include <sys/time.h>
#include "deconv.h"
#include "utils.h"
#include "unetv1.h"
int main()
{
    // Test parameters
    std::array<uint32_t, 2> ifm_shape = {N_IM, N_ICH};
    std::array<uint32_t, 2> ofm_shape = {N_OM, N_OCH};
    data_stream<BIT_ACTV> s_input("s_input");
    data_stream<BIT_ACTV> s_input_golden("s_input_golden");
    gen_ifm(s_input, ifm_shape);
    gen_ifm(s_input_golden, ifm_shape);

    ap_int<BIT_WGHT> weight[N_OCH][K*K][N_ICH];
    for(uint32_t ic = 0; ic < N_ICH; ++ic) {
        for(uint32_t oc = 0; oc < N_OCH; ++oc) {
            for(uint32_t k = 0; k < K*K; ++k) {
                ap_int<BIT_WGHT> val = ic*N_OCH*K*K + oc*K*K + k;
                weight[oc][k][ic] = val;
            }
        }
    }

    for (uint32_t fo = 0; fo < FOLD_O; ++fo) {
        for (uint32_t fi = 0; fi < FOLD_I; ++fi) {
            for (uint32_t k = 0; k < K*K; ++k) {
                ap_uint<P_OCH * P_ICH * BIT_WGHT> w = 0;
                for (uint32_t o = 0; o < P_OCH; ++o) {
                    for (uint32_t i = 0; i < P_ICH; ++i) {
                        uint32_t oc = fo * P_OCH + o;
                        uint32_t ic = fi * P_ICH + i;
                        ap_uint<BIT_WGHT> val = weight[oc][k][ic];
                        w(SLICE(BIT_WGHT, P_ICH * o + i)) = val;
                    }
                }
                weight_reshape[fo][fi][k] = w;
            }
        }
    }
    axiu_stream<BIT_CONV> s_output_axiu("s_output_axiu");
    unetv1(s_input, s_output_axiu);

    data_stream<BIT_CONV> s_output("s_output");
    while(s_output_axiu.size() > 0) {
        auto temp = s_output_axiu.read();
        s_output.write(temp.data);
    }

    data_stream<BIT_CONV> s_output_golden("s_output_golden");
    deconv_golden<N_ICH, N_OCH, N_IH, N_IW, K, P, S, O_P, BIT_ACTV, BIT_WGHT, BIT_CONV>(
        s_input_golden, s_output_golden, weight
    );
    auto error_cnt = check_afm<BIT_CONV, 2>(s_output_golden, s_output, ofm_shape, 1);
    if(error_cnt > 0) {
        std::cout << "Test Failed! " << error_cnt << " errors" << std::endl;
        return 1;
    } else {
        std::cout << "Test Passed! " << std::endl;
        return 0;
    }
}